package com.idanchuang.component.mq.pulsar;

import com.alibaba.fastjson.JSON;
import com.idanchuang.base.spring.intercept.MessageInterceptorCallback;
import com.idanchuang.component.core.helper.EnvContext;
import org.apache.pulsar.client.api.Message;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.URLDecoder;
import java.net.URLEncoder;
import java.util.Map;

/**
 * @author yjy
 * Created at 2022/2/28 5:23 下午
 */
public class PressureMessageInterceptorCallback implements MessageInterceptorCallback {

    private static final Logger logger = LoggerFactory.getLogger(PressureMessageInterceptorCallback.class);

    private final ComponentPulsarConfigBean componentPulsarConfigBean;

    public PressureMessageInterceptorCallback(ComponentPulsarConfigBean componentPulsarConfigBean) {
        this.componentPulsarConfigBean = componentPulsarConfigBean;
    }

    @Override
    public void beforSend(String topic, Object message, Map<String, String> properties) {
        try {
            if (!componentPulsarConfigBean.isEnabled()) {
                return;
            }
            // 如果有环境变量, 则传递
            if (EnvContext.getAttributeCount() > 0) {
                String attributes = JSON.toJSONString(EnvContext.getAttributes());
                attributes = URLEncoder.encode(attributes, "UTF-8");
                properties.put(EnvContext.TRANSMIT_ATTRIBUTE_HEADER_NAME, attributes);
                logger.debug("beforSend add properties: {}, value: {}", EnvContext.TRANSMIT_ATTRIBUTE_HEADER_NAME, attributes);
            }
        } catch (Exception e) {
            logger.warn("beforSend error", e);
        }
    }

    @Override
    public void beforeConsume(Message message) {
       try {
           if (!componentPulsarConfigBean.isEnabled()) {
               return;
           }
           // 清除当前环境变量
           EnvContext.clearContext();
           Object attributes = message.getProperties().get(EnvContext.TRANSMIT_ATTRIBUTE_HEADER_NAME);
           if (attributes != null) {
               String attributesStr = URLDecoder.decode(attributes.toString(), "UTF-8");
               EnvContext.setAllAttribute(JSON.parseObject(attributesStr));
               logger.debug("beforeConsume get properties: {}, value: {}", EnvContext.TRANSMIT_ATTRIBUTE_HEADER_NAME, attributesStr);
           }
           EnvContext.setDcTraceIdIfAbsent();
       } catch (Exception e) {
           logger.warn("beforeConsume error", e);
       }
    }


}
